import pickle
from sentence_transformers import SentenceTransformer
import os
import ssl
import requests
from huggingface_hub import snapshot_download

# 加载预处理数据
def load_preprocessed_data(file_path):
    with open(file_path, 'rb') as f:
        return pickle.load(f)

# 生成文本嵌入
def generate_embeddings(texts):
    try:
        # 首先尝试离线加载
        model = SentenceTransformer('all-MiniLM-L6-v2', device='cpu')
    except Exception as e:
        print("正在尝试下载模型...")
        try:
            # 创建自定义 SSL 上下文
            ssl_context = ssl.create_default_context()
            ssl_context.check_hostname = False
            ssl_context.verify_mode = ssl.CERT_NONE
            
            # 下载模型
            model_path = snapshot_download(
                repo_id="sentence-transformers/all-MiniLM-L6-v2",
                local_files_only=False,
                ssl_verify=False
            )
            model = SentenceTransformer(model_path, device='cpu')
        except Exception as download_error:
            print(f"模型下载失败: {download_error}")
            raise

    print("模型加载成功，开始生成嵌入...")
    return model.encode(texts, show_progress_bar=True)

# 保存嵌入
def save_embeddings(embeddings, output_path):
    with open(output_path, 'wb') as f:
        pickle.dump(embeddings, f)

def main():
    preprocessed_data_path = "data/preprocessed_data.pkl"
    embeddings_output_path = "embeddings/knowledge_embeddings.pkl"
    
    # 确保 embeddings 文件夹存在
    os.makedirs("embeddings", exist_ok=True)
    
    print("正在加载预处理数据...")
    preprocessed_data = load_preprocessed_data(preprocessed_data_path)
    
    print(f"开始生成嵌入，共 {len(preprocessed_data)} 条数据...")
    embeddings = generate_embeddings(preprocessed_data)
    
    print("正在保存嵌入...")
    save_embeddings(embeddings, embeddings_output_path)
    print(f"嵌入已保存到 {embeddings_output_path}")

if __name__ == "__main__":
    main()
